# BSD3.0
# cili@163.com

import socketserver
import struct

class DNSQuery:
    def __init__(self, data):
        for i in range(len(data)):
            d = data[i + 1]
            if d == 0: break
            if d < 32:
                self.name = self.name + '.'
            else:
                self.name = self.name + chr(d)

        self.querybytes = data[0:i + 2]
        (self.type, self.classify) = struct.unpack('>HH', data[i + 2:i + 6])
        self.len = i + 6
    def getbytes(self):
        return self.querybytes + struct.pack('>HH', self.type, self.classify)


class DNSAnswer:
    def __init__(self, ip):
        self.name = 49164
        self.type = 1
        self.classify = 1
        self.timetolive = 300
        self.datalength = 4
        self.ip = ip

    def getbytes(self):
        res = struct.pack('>HHHLH', self.name, self.type, self.classify, self.timetolive, self.datalength)
        s = self.ip.split('.')
        res = res + struct.pack('BBBB', int(s[0]), int(s[1]), int(s[2]), int(s[3]))
        return res

class DNSFrame:
    def __init__(self, data):
        (self.id, self.flags, self.quests, self.answers, self.author, self.addition) = struct.unpack('>HHHHHH',
                                                                                                     data[0:12])
        self.query = DNSQuery(data[12:])

    def setip(self, ip):
        self.answer = DNSAnswer(ip)
        self.answers = 1
        self.flags = 33152

    def getbytes(self):
        res = struct.pack('>HHHHHH', self.id, self.flags, self.quests, self.answers, self.author, self.addition)
        res = res + self.query.getbytes()
        if self.answers != 0: res = res + self.answer.getbytes()
        return res


class DNSUDPHandler(socketserver.BaseRequestHandler):
    def handle(self):
        data = self.request[0].strip()
        dns = DNSFrame(data)
        socket = self.request[1]
        namemap = DNSServer.namemap
        if (dns.query.type == 1):
            name = dns.query.name;
            if name in namemap:
                dns.setip(namemap[name])
                socket.sendto(dns.getbytes(), self.client_address)
        else:
            socket.sendto(data, self.client_address)


class DNSServer:
    namemap = {}
    def __init__(self, port=53):
        self.port = port

    def add(self, name, ip):
        DNSServer.namemap[name] = ip

    def update(self,data):
        DNSServer.namemap.clear()
        for name,ip in data:
            DNSServer.namemap[name] = ip
        DNSServer.namemap['*']='0.0.0.0'

    def start(self):
        server = socketserver.UDPServer(("0.0.0.0", self.port), DNSUDPHandler)
        server.serve_forever()

    def stop(self):
        if self.server:
            self.server.shutdown()
